import torch
import clip
from PIL import Image
import torch.nn.functional as F

from clip.model import build_model
from clip.clip import _transform


if __name__ == "__main__":
    checkpoint_path = r"E:\checkpoint\open_clip_laion_2B\open_clip_pytorch_model.bin"
    image_path = r"E:\BL_dataset\舆情数据集\xinyu\datasets\126.jpg"
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    image = Image.open(image_path)

    state_dict = torch.load(checkpoint_path)
    model = build_model(state_dict).cuda()

    image_processor = _transform(model.visual.input_resolution)
    text_tokenizer = clip.tokenize

    image = image_processor(image).unsqueeze(0).to(device)
    text = text_tokenizer(["双剑老狗，大家听说过吗？"]).to(device)

    with torch.no_grad():
        image_features = model.encode_image(image)
        text_features = model.encode_text(text)

    # 归一化
    image_features /= image_features.norm(dim=-1, keepdim=True)
    text_features /= text_features.norm(dim=-1, keepdim=True)

    similarity = (100.0 * image_features @ text_features.T).cpu().float()
    print(similarity.softmax(dim=-1).tolist())


